import os
import base64
import requests

from io import BytesIO
from PIL import Image, ImageOps



""" 以下是通过本地模型获取图片已经对应信息的函数 """
def process_date(PIPELINE, input_dir="/root/autodl-tmp/date"):

    all_data = []
    for filename in os.listdir(input_dir):
        if filename.lower().endswith((".jpg", ".jpeg", ".png", ".bmp")):
            image_path = os.path.join(input_dir, filename)
            output = PIPELINE.predict(image_path)
            for res in output:
                all_data.append(res)
    
    return all_data


def optimize_seal_image(input_path, min_size=512, border=30):
    with Image.open(input_path) as img:
        w, h = img.size
        # 放大图像（短边>=512）
        if min(w, h) < min_size:
            scale = min_size / min(w, h)
            new_size = (int(w * scale), int(h * scale))
            img = img.resize(new_size, Image.LANCZOS)
        # 加边框增强上下文
        img = ImageOps.expand(img, border=border, fill='white')
        # 保存临时优化图
        optimized_path = os.path.join("/tmp", os.path.basename(input_path))
        img.save(optimized_path)
        return optimized_path


def process_stamp(PIPELINE_SEAL, input_dir="/root/autodl-tmp/stamp"):
    all_stamp = []
    for filename in os.listdir(input_dir):
        if filename.lower().endswith((".jpg", ".jpeg", ".png", ".bmp")):
            image_path = os.path.join(input_dir, filename)
            print(f"📌 正在处理：{image_path}")

            optimized_img = optimize_seal_image(image_path)
            output = PIPELINE_SEAL.predict(
                optimized_img,
                use_doc_orientation_classify=True,
                use_doc_unwarping=True,
            )

            for res in output:
                all_stamp.append(res)

    return all_stamp


def process_handwrite(PIPELINE, input_dir="/root/autodl-tmp/handwrite"):
    # pipeline = create_model(model_name="PP-OCRv5_server_rec")

    all_handwrite = []
    for filename in os.listdir(input_dir):
        if filename.lower().endswith((".jpg", ".jpeg", ".png", ".bmp")):
            image_path = os.path.join(input_dir, filename)
            output = PIPELINE.predict(image_path)
            for res in output:
                all_handwrite.append(res)

    return all_handwrite


def process_ocr(PIPELINE_OCR, img_base64):
    import numpy as np
    all_ocr = []
    
    image_data = base64.b64decode(img_base64)
    img = Image.open(BytesIO(image_data)).convert('RGB')
    img_array = np.array(img)

    output = PIPELINE_OCR.predict(
        input=img_array,
        use_doc_orientation_classify=False,
        use_doc_unwarping=False,
        use_textline_orientation=False,
    )
    for res in output:
        all_ocr.append(res)

    return all_ocr


def process_page(PIPELINE, input_dir = "/root/autodl-tmp/pagenum"):
    # pipeline = create_model(model_name="PP-OCRv5_server_rec")

    all_pagenum = []
    for filename in os.listdir(input_dir):
        if filename.lower().endswith((".jpg", ".jpeg", ".png", ".bmp")):
            image_path = os.path.join(input_dir, filename)
            output = PIPELINE.predict(image_path)

            # 处理结果输出
            for res in output:
                all_pagenum.append(res)

    return all_pagenum





""" 以下是通过api获取图片已经对应信息的函数 """
def process_image(img_base64):
    # 创建保存目录
    os.makedirs('stamp', exist_ok=True)
    os.makedirs('handwrite', exist_ok=True)
    os.makedirs('date', exist_ok=True)

    api_url = "http://183.221.0.158:29986/api/algo/docLayout"  #公网ip

    try:
        response = requests.post(api_url, json={"img_base64": img_base64})
        print(f"HTTP状态码: {response.status_code}")  # 调试信息
        response.raise_for_status()
    except Exception as e:
        print(f"API请求失败: {str(e)}")
        return

    print("原始响应内容:", response.text[:200])  # 截取前200字符避免刷屏
    
    try:
        result = response.json()
    except Exception as e:
        print(f"JSON解析失败: {str(e)}")
        return
    
    print("解析后的数据结构:", result)  # 显示完整数据结构

    if result.get('code', 1) != 0:  # 默认非零表示错误
        print(f"接口业务错误: [{result.get('code')}] {result.get('msg')}")
        return

    if not isinstance(result.get('data'), list):
        print("无效的data字段结构")
        return

    try:
        image_data = base64.b64decode(img_base64)
        img = Image.open(BytesIO(image_data)).convert('RGB')
    except Exception as e:
        print(f"图片处理失败: {str(e)}")
        return

    for i, item in enumerate(result['data']):
        try:
            if not all(key in item for key in ['position', 'category']):
                print(f"区域{i}缺少关键字段")
                continue
                
            x, y, w, h = map(int, item['position'])
            region = img.crop((x, y, x + w, y + h))
            
            category = item['category'].strip().lower()
            save_dir = {'stamp': 'stamp', 'handwrite': 'handwrite', 'date': 'date'}.get(category, 'other')
            
            filename = f"{save_dir}/region_{i+1}_{os.urandom(4).hex()}.jpg"
            region.save(filename)
            print(f"成功保存: {filename}")
            
        except Exception as e:
            print(f"区域{i}处理异常: {str(e)}")

    return result



def get_page_number(img_base64):
    # 创建保存目录
    os.makedirs('pagenum', exist_ok=True)

    api_url = "http://183.221.0.158:29986/api/algo/pageNum"

    try:
        response = requests.post(api_url, json={"img_base64": img_base64})
        print(f"HTTP状态码: {response.status_code}")  # 调试信息
        response.raise_for_status()
    except Exception as e:
        print(f"API请求失败: {str(e)}")
        return

    print("原始响应内容:", response.text[:200])  # 截取前200字符避免刷屏
    
    try:
        result = response.json()
    except Exception as e:
        print(f"JSON解析失败: {str(e)}")
        return
    
    print("解析后的数据结构:", result)

    if result.get('code', 1) != 0:
        print(f" PageNumber 接口业务错误: [{result.get('code')}] {result.get('msg')}")
        return

    if not isinstance(result.get('data'), list):
        print("无效的data字段结构")
        return

    try:
        image_data = base64.b64decode(img_base64)
        img = Image.open(BytesIO(image_data)).convert('RGB')
    except Exception as e:
        print(f"图片处理失败: {str(e)}")
        return

    for i, item in enumerate(result['data']):
        try:
            if not all(key in item for key in ['position', 'category']):
                print(f"区域{i}缺少关键字段")
                continue

            x, y, w, h = map(int, item['position'])
            region = img.crop((x, y, x + w, y + h))

            category = item['category'].strip().lower()
            if category != "pagenum":
                print(f"区域{i}不是页码，跳过")
                continue

            filename = f"pagenum/page_{i+1}_{os.urandom(4).hex()}.jpg"
            region.save(filename)
            print(f"成功保存页码图像: {filename}")

        except Exception as e:
            print(f"区域{i}处理异常: {str(e)}")

    return result

